import time

import numpy
import numpy as np
import matplotlib

import particle_tracer.particle_generator

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
from _logging import logger
plt.ion()


# 常数
epsilon_0 = 8.85e-12  # 真空介电常数 (F/m)
# Q = 1e-9  # 电荷量 (C)
Q = 2e-7 * 50

L = 0.2  # 计算区域长度 (单位：m)
N = 200  # 网格点数

# 网格
z = np.linspace(-L/2, L/2, N)  # 离散的z轴坐标

# 电荷位置和带电量
def charge_dens(z):
    return numpy.exp(-z**2 / (2* 0.01 **2))
z_initial = particle_tracer.particle_generator.ParticlePositionGenerator(charge_dens).generate(200,-0.04,0.04,1e-5)#np.linspace(-0.1, 0.1, 50)

qi = Q / len(z_initial)
dz = z[1] - z[0]  # 网格间距
plt.figure()
plt.plot(z_initial,numpy.zeros(z_initial.shape),'.')


plt.figure()
plt.plot(z, charge_dens(z))
charge_positions =z_initial  # 电荷位置

# 构建电荷密度分布
rho = np.zeros(N)  # 初始电荷密度
for z_i in charge_positions:
    # 采用狄拉克δ函数近似为高斯函数
    rho += qi * np.exp(-(z - z_i)**2 / (2 * dz**2)) / (np.sqrt(2 * np.pi) * dz)
# plt.figure()
# plt.plot(z,rho)
# 计算电势：使用高斯积分和有限差分法解泊松方程
phi = np.zeros(N)


# 二阶差分公式：d^2(phi)/dz^2 ≈ (phi[i+1] - 2*phi[i] + phi[i-1]) / dz^2
# 设置边界条件：远离电荷时电势趋近于零
phi[0] = 0
phi[-1] = 0

logger.info("直接矩阵求逆计算电势")

t1 = time.time()
M = []
for i in range(N-1):
    m = numpy.zeros(N+1 )
    m[i:i+3] = [-1/2,1,-1/2]
    M.append(m)
M = numpy.matrix(M)[:-1,1:-2]
# phi[1:-1] = numpy.linalg.pinv
phi[1:-1] =( M**-1* numpy.matrix(rho[1:-1]).T).A.ravel() *dz**2 / (2 * epsilon_0)
logger.info("Time cost = %.5f s"%(time.time() - t1))
# 计算电场：E = -d(phi)/dz
E = -np.gradient(phi, dz)

# 绘制电荷密度和电场分布
# 电荷密度分布
plt.figure(constrained_layout = True)
plt.plot(z, rho, label=r'$\rho(z)$', color='b')
plt.xlabel('Position (z) [m]')
plt.ylabel('Charge Density (C/m)')

# 电场分布
plt.figure(constrained_layout = True)
plt.plot(z, E, label=r'$E(z)$', color='r')
plt.xlabel('Position (z) [m]')
plt.ylabel('Electric Field (N/C)')


logger.info("迭代计算电势：使用高斯-塞德尔迭代法")


phi_history = []


# 迭代计算电势：使用高斯-塞德尔迭代法
t2 = time.time()
iters = list(range(10000))
for iter in iters:  # 设置迭代次数
    for i in range(1, N-1):
        phi[i] = 0.5 * (phi[i-1] + phi[i+1] + (rho[i] * dz**2) / epsilon_0)
    phi_history.append(phi[len(phi) //2])
logger.info("Time cost = %.5f s"%(time.time() - t2))

plt.figure()
plt.plot(iters, phi_history)
plt.xlabel("iteration")
plt.ylabel("phi(x = mid)")

# 计算电场：E = -d(phi)/dz
E = -np.gradient(phi, dz)

# 绘制电荷密度和电场分布
# 电荷密度分布
plt.figure(constrained_layout = True)
plt.plot(z, rho, label=r'$\rho(z)$', color='b')
plt.xlabel('Position (z) [m]')
plt.ylabel('Charge Density (C/m)')

# 电场分布
plt.figure(constrained_layout = True)
plt.plot(z, E, label=r'$E(z)$', color='r')
plt.xlabel('Position (z) [m]')
plt.ylabel('Electric Field (N/C)')

plt.show()
